#!/usr/bin/env python3
"""
Script to fix incorrect \chapter commands inside LaTeX code listings.
These should be converted to comments or proper section markers.
"""

import re
import glob
import os
from pathlib import Path

def fix_chapter_commands_in_file(file_path):
    """Fix \chapter commands that appear in code listings in a LaTeX file."""
    print(f"Processing: {file_path}")
    
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    original_content = content
    changes_made = 0
    
    # Pattern 1: Fix \chapter commands that appear after \end{lstlisting}[language]
    # These are actually part of code blocks and should be comments
    patterns_to_fix = [
        # Pattern: \end{lstlisting}python\n\chapter{...}
        (r'\\end\{lstlisting\}(python|bash|dockerfile|toml|yaml|json)\s*\n\\chapter\{([^}]+)\}',
         r'\\end{lstlisting}\n\n\\begin{lstlisting}[language=\1]\n# \2'),
        
        # Pattern: standalone \chapter{...} inside code blocks (convert to comments)
        # This is more complex - we need to detect when we're inside a code block
    ]
    
    for pattern, replacement in patterns_to_fix:
        new_content = re.sub(pattern, replacement, content, flags=re.MULTILINE)
        if new_content != content:
            changes_count = len(re.findall(pattern, content, flags=re.MULTILINE))
            changes_made += changes_count
            print(f"  - Fixed {changes_count} instances of pattern: {pattern[:50]}...")
            content = new_content
    
    # More sophisticated approach: track when we're inside lstlisting blocks
    lines = content.split('\n')
    in_listing = False
    listing_language = None
    fixed_lines = []
    line_changes = 0
    
    for i, line in enumerate(lines):
        # Check if we're entering a listing
        begin_match = re.match(r'\\begin\{lstlisting\}(?:\[language=(\w+)\])?', line.strip())
        if begin_match:
            in_listing = True
            listing_language = begin_match.group(1) if begin_match.group(1) else None
            fixed_lines.append(line)
            continue
        
        # Check if we're exiting a listing
        if re.match(r'\\end\{lstlisting\}', line.strip()):
            in_listing = False
            listing_language = None
            fixed_lines.append(line)
            continue
        
        # If we're inside a listing and find a \chapter command, convert it to a comment
        if in_listing and line.strip().startswith('\\chapter{'):
            chapter_match = re.match(r'\\chapter\{([^}]+)\}', line.strip())
            if chapter_match:
                chapter_text = chapter_match.group(1)
                # Convert to appropriate comment style based on language
                if listing_language in ['python', 'bash', 'yaml']:
                    fixed_lines.append(f"# {chapter_text}")
                elif listing_language in ['dockerfile']:
                    fixed_lines.append(f"# {chapter_text}")
                elif listing_language in ['toml']:
                    fixed_lines.append(f"# {chapter_text}")
                else:
                    # Default to hash comment
                    fixed_lines.append(f"# {chapter_text}")
                line_changes += 1
                print(f"  - Line {i+1}: Converted \\chapter{{{chapter_text}}} to comment")
                continue
        
        fixed_lines.append(line)
    
    if line_changes > 0:
        content = '\n'.join(fixed_lines)
        changes_made += line_changes
    
    # Additional pattern fixes for specific cases
    additional_fixes = [
        # Fix cases where \chapter appears right after code without proper separation
        (r'([\w\}\]\)\;])\s*\n\\chapter\{([^}]+)\}\n([a-zA-Z_])', r'\1\n\n# \2\n\3'),
        
        # Fix \subsubsection inside lstlisting blocks
        (r'\\end\{lstlisting\}\s*\n\\subsubsection\{([^}]+)\}\s*\n\\end\{lstlisting\}([a-z]+)', 
         r'\\end{lstlisting}\n\n\\subsubsection{\1}\n\n\\begin{lstlisting}[language=\2]'),
    ]
    
    for pattern, replacement in additional_fixes:
        new_content = re.sub(pattern, replacement, content, flags=re.MULTILINE)
        if new_content != content:
            fix_count = len(re.findall(pattern, content, flags=re.MULTILINE))
            changes_made += fix_count
            print(f"  - Applied additional fix: {fix_count} instances")
            content = new_content
    
    # Write the file back if changes were made
    if content != original_content:
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(content)
        print(f"  ✓ Saved {changes_made} changes to {file_path}")
        return changes_made
    else:
        print(f"  - No changes needed for {file_path}")
        return 0

def main():
    """Process all .tex files in the chapters directory."""
    print("LaTeX Chapter Command Fixer")
    print("=" * 50)
    
    # Find all .tex files in the chapters directory
    chapters_dir = Path("chapters")
    if not chapters_dir.exists():
        print("Error: 'chapters' directory not found!")
        return
    
    tex_files = list(chapters_dir.glob("*.tex"))
    if not tex_files:
        print("No .tex files found in chapters directory!")
        return
    
    print(f"Found {len(tex_files)} LaTeX files to process:")
    for f in tex_files:
        print(f"  - {f}")
    print()
    
    total_changes = 0
    processed_files = 0
    
    for tex_file in tex_files:
        try:
            changes = fix_chapter_commands_in_file(tex_file)
            total_changes += changes
            processed_files += 1
        except Exception as e:
            print(f"Error processing {tex_file}: {e}")
    
    print("\n" + "=" * 50)
    print(f"Processing complete!")
    print(f"Files processed: {processed_files}/{len(tex_files)}")
    print(f"Total changes made: {total_changes}")
    
    if total_changes > 0:
        print("\nRecommendation: Review the changes and test LaTeX compilation.")

if __name__ == "__main__":
    main()